package com.zyh.langchain4j.config;

import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.bind.annotation.RestControllerAdvice;

import java.util.List;

@Configuration
@RestControllerAdvice
public class ConsultantConfig {
    private final OpenAiChatModel openAiChatModel;
    private final ChatMemoryStore redisChatMemoryStore;
    // 向量模型
    private final EmbeddingModel embeddingModel;
    // 向量数据库
    private final RedisEmbeddingStore redisEmbeddingStore;

    public ConsultantConfig(OpenAiChatModel openAiChatModel, ChatMemoryStore redisChatMemoryStore, EmbeddingModel embeddingModel, RedisEmbeddingStore redisEmbeddingStore) {
        this.openAiChatModel = openAiChatModel;
        this.redisChatMemoryStore = redisChatMemoryStore;
        this.embeddingModel = embeddingModel;
        this.redisEmbeddingStore = redisEmbeddingStore;
    }

    /*    *//**
     * 创建代理对象
     *
     * @return
     *//*
    @Bean
    public ConsultantService consultantService() {
        ConsultantService service = AiServices.builder(ConsultantService.class)
                .chatModel(openAiChatModel)
                .build();
        return service;
    }*/

    /**
     * 创建会话记忆对象
     *
     * @return
     */
    @Bean
    public ChatMemory chatMemory() {
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .maxMessages(20)//最大信息数量
                .build();
        return chatMemory;
    }

    /**
     * 创建会话记忆隔离对象ChatMemoryProvider
     *
     * @return
     */
    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
        ChatMemoryProvider chatMemoryProvider = new ChatMemoryProvider() {
            @Override
            public ChatMemory get(Object memoryId) {
                return MessageWindowChatMemory.builder()
                        .id(memoryId)//设置记忆id(数据隔离)
                        .maxMessages(20)//最大信息数量
                        .chatMemoryStore(redisChatMemoryStore)//使用自定义的会话记忆存储
                        .build();
            }
        };
        return chatMemoryProvider;
    }

    /**
     * 创建向量数据库存储对象
     *
     * @return
     */
//    @Bean
    public EmbeddingStore store() {
        //1.加载文件进内存（默认解析器）
        List<Document> contents = ClassPathDocumentLoader.loadDocuments("content");
        // 专门解析pdf文件，
//        List<Document> contents = ClassPathDocumentLoader.loadDocuments("content",new ApachePdfBoxDocumentParser());
//        List<Document> contents = FileSystemDocumentLoader.loadDocuments("E:\\IDEA\\AI\\LangChain4j\\LangChain4j\\src\\main\\resources\\content");
//        List<Document> contents = UrlDocumentLoader.load("");
        //2.构建向量数据库操作对象
        // 2.1 内存存储
//        InMemoryEmbeddingStore store = new InMemoryEmbeddingStore();
        // 2.2 redis向量数据库存储

        // 3.构建文本分割器对象(最大分割字符数量，上一篇段的尾部取多少和下一片段合并)
        DocumentSplitter splitter = DocumentSplitters.recursive(500, 100);
        //4.构造，完成文本数据切割，向量化然后存储
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
//                .embeddingStore(store)
                .embeddingStore(redisEmbeddingStore)
                .documentSplitter(splitter)//文本切割器
                .embeddingModel(embeddingModel)//向量化模型
                .build();
        ingestor.ingest(contents);
        return redisEmbeddingStore;
    }

    /**
     * 创建数据库内容检索对象
     *
     * @param store
     * @return
     */
    @Bean
    public ContentRetriever contentRetriever(EmbeddingStore store) {
        return EmbeddingStoreContentRetriever.builder()
//                .embeddingStore(store)//向量数据库
                .embeddingStore(redisEmbeddingStore)//向量数据库
                .maxResults(3)//最多返回片段
                .minScore(0.7)//匹配最低向量值
                .embeddingModel(embeddingModel)//向量化模型
                .build();
    }
}
